import torch
from torch import nn
from torch.nn import TransformerEncoderLayer


class PositionalEncoding(nn.Module):
    # ...existing code from original Bert.py...
    def __init__(self, d_model, max_len=100):
        super(PositionalEncoding, self).__init__()
        self.d_model = d_model
        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-torch.log(torch.tensor(10000.0)) / d_model))
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        pe = pe.unsqueeze(0).transpose(0, 1)
        self.register_buffer('pe', pe)

    def forward(self, x):
        # x: (seq_len, batch_size, d_model)
        x = x + self.pe[:x.size(0), :]
        return x


class PositionEncoder(nn.Module):
    # ...existing code from original Bert.py...
    def __init__(self, input_dim=2, d_model=128, nhead=8, num_layers=2, max_nodes=100):
        super(PositionEncoder, self).__init__()
        self.d_model = d_model
        self.input_projection = nn.Linear(input_dim, d_model)
        self.pos_encoding = PositionalEncoding(d_model, max_len=max_nodes)
        encoder_layer = TransformerEncoderLayer(d_model, nhead)
        self.transformer_encoder = nn.TransformerEncoder(encoder_layer, num_layers)
        self.output_projection = nn.Linear(d_model, 64)

    def forward(self, pos_features, graph_batch=None):
        x = self.input_projection(pos_features)
        x = self.transformer_encoder(x)
        x = self.output_projection(x.squeeze(1))
        return x

